/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include <ntddk.h>
#include "cpu.h"
#include "vmx.h"
#include "vt.h"
#include "io.h"
#include "gui.h"
#include "port.h"

CHAR gScanCodeMap[255];

NTSTATUS 
DescIoExitQualification(
		IN ULONG ExitQualification,
		OUT PIO_PORT_INFO PortInfo
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	if ( PortInfo == NULL ){
		return STATUS_UNSUCCESSFUL;
	}

	PortInfo->Port = (ExitQualification & 0xFFFF0000) >> 16;
	PortInfo->Size = (ExitQualification & 7) + 1;
	PortInfo->State = (ExitQualification & (1 << 3)) ? IO_IN : IO_OUT;
	PortInfo->Type = (ExitQualification & (1 << 4)) ? IO_STRING : IO_CHAR;
	PortInfo->Rep = (ExitQualification & (1 << 5)) ? IO_REP : IO_OTHER;
	
	return STATUS_SUCCESS;
}

#define CPU_BASED_PRIMARY_IO 25
VOID
SetHandleIOAccess(
      BOOLEAN IsSet
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG v;

    v = _ReadVMCS(CPU_BASED_VM_EXEC_CONTROL);

    if (IsSet) {
        /* Enable I/O bitmaps */
        CmSetBit32(&v,   CPU_BASED_PRIMARY_IO);
    } else {
        /* Disable I/O bitmaps */
        CmClearBit32(&v, CPU_BASED_PRIMARY_IO);
    }

    _WriteVMCS(CPU_BASED_VM_EXEC_CONTROL, v);
}
VOID 
HandleIoAccess(
        PCPU_VM_CONTEXT pCPU,
        PSAVED_REGISTOR pSavedReg
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG InstructionLength;
    ULONG ExitQualification;
    UCHAR Char;
    BOOLEAN bMouse;

	IO_PORT_INFO IoPortInfo;
	
	
	//
	// Get the ExitQualification Info
	//
	
	ExitQualification = _ReadVMCS(EXIT_QUALIFICATION);
	DescIoExitQualification( ExitQualification, &IoPortInfo );
	
	if (IoPortInfo.Size != 1 || IoPortInfo.Type == IO_STRING || \
		IoPortInfo.Rep == IO_REP) {
		goto GoResume;
	}
	
	if (KeyboardReadKeystroke(&Char, FALSE, \
		&bMouse) != STATUS_SUCCESS) {
		goto GoResume;
	}
	
	pSavedReg->REAX = ((ULONG)pSavedReg->REAX & 0xFFFFFF00) | (Char & 0xFF);
	
	if (!bMouse && Char == 88) {
		KdPrint(("!!!!F12 Press!!!!!\n"));
		
		//
		// Enter Debuger GUI
		//
		
		GUIEntryDebug(pCPU, pSavedReg);
        return;
	}

GoResume:

	//
	// Go Next
	//
	
    InstructionLength = _ReadVMCS(VM_EXIT_INSTRUCTION_LEN);
    _WriteVMCS(GUEST_RIP, _ReadVMCS(GUEST_RIP)+InstructionLength);

}

NTSTATUS
i8042ReadKeyboardData(
		UCHAR* pChar, 
		PBOOLEAN bMouse
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	UCHAR PortStatus;

	PortStatus = IoReadPortByte(KEYB_REGISTER_STATUS);

	if (PortStatus & KEYB_STATUS_OBUFFER_FULL) {
		
		//
		// Data is available 
		//

		*pChar = IoReadPortByte(KEYB_REGISTER_DATA);
		if ((PortStatus & KEYB_STATUS_PARITY_ERROR) == 0) {
			
			//
			// Check if this is a mouse event or not
			//

			*bMouse = (PortStatus & KEYB_STATUS_TRANSMIT_TIMEOUT) != 0;
			return STATUS_SUCCESS;
		}
	}

	return STATUS_UNSUCCESSFUL;
}

BOOLEAN 
i8042WriteKeyboardData(
		USHORT Port, 
		UCHAR Data
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG Counter;

	Counter = POLL_STATUS_ITERATIONS;
	while (	(KEYB_STATUS_IBUFFER_FULL & \
			 IoReadPortByte(KEYB_REGISTER_STATUS)) && \
		    (Counter--)) {
			KeStallExecutionProcessor(1);
	}

	if (Counter) {
		IoWritePortByte(Port, Data);
		return TRUE;
	}

	return FALSE;
}

NTSTATUS
KeyboardReadKeystroke(
		PUCHAR pChar, 
		BOOLEAN UnGet, 
		PBOOLEAN bMouse
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG Counter;
	UCHAR PortStatus, ScanCode;
	NTSTATUS Status;

	Counter = POLL_STATUS_ITERATIONS;
	while (Counter) {
		PortStatus = IoReadPortByte(KEYB_REGISTER_STATUS);

		Status = i8042ReadKeyboardData(&ScanCode, bMouse);

		if (Status == STATUS_SUCCESS) {
			break;
		}

		KeStallExecutionProcessor(1);
		Counter--;
	}

	if (Counter == 0) {
		return STATUS_UNSUCCESSFUL;
	}

	if (UnGet) {

		//
		// Echo back the scancode
		//

		i8042WriteKeyboardData(KEYB_REGISTER_COMMAND, KEYB_COMMAND_DISABLE_KEYBOARD);
		i8042WriteKeyboardData(KEYB_REGISTER_COMMAND, KEYB_COMMAND_WRITE_OUTPUT);
		i8042WriteKeyboardData(KEYB_REGISTER_DATA, ScanCode);
		i8042WriteKeyboardData(KEYB_REGISTER_COMMAND, KEYB_COMMAND_ENABLE_KEYBOARD);
	}

	*pChar = ScanCode;

	return STATUS_SUCCESS;
}

BOOLEAN 
ChariIsLower(
	CHAR c
)
{
	return (c >= 'a' && c <= 'z');
}

UCHAR 
CharToUpper(
	UCHAR c
)
{
	if (!ChariIsLower(c))
		return c;

	return (c - 'a' + 'A');
}

KEYBOARD_STATUS gKeyboardStatus={0};

UCHAR 
KeyboardScancodeToKeycode(
		UCHAR Char
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	BOOLEAN Handled;

	Handled = FALSE;

	switch(Char & ~SCANCODE_RELEASE_FLAG) {
	case 0x1d:
		gKeyboardStatus.lctrl = IS_SCANCODE_RELEASE(Char) ? FALSE : TRUE;
		Handled = TRUE;
		break;
	case 0x2a:
		gKeyboardStatus.lshift = IS_SCANCODE_RELEASE(Char) ? FALSE : TRUE;
		Handled = TRUE;
		break;
	case 0x36:
		gKeyboardStatus.rshift = IS_SCANCODE_RELEASE(Char) ? FALSE : TRUE;
		Handled = TRUE;
		break;
	case 0x38:
		gKeyboardStatus.lalt = IS_SCANCODE_RELEASE(Char) ? FALSE : TRUE;
		Handled = TRUE;
		break;
	default:
		break;
	}

	if (Handled)
		return 0;
	
	//
	// No need to do anything else for every other released key
	//

	if IS_SCANCODE_RELEASE(Char)
		return 0;

	//
	// Else, ignore errors and acks, let the guest OS handle them if needed
	//

	switch(Char) {
		case 0x00: /* KBD ERROR -> ignore */
		case 0xaa: /* BAT OK -> ignore */
		case 0xee: /* ECHO CMD RES -> ignore */
		case 0xfa: /* ACK FROM KBD -> ignore */
		case 0xfc: /* BAT ERROR -> ignore? */
		case 0xfd: /* INTERNAL FAILURE -> ignore? */
		case 0xfe: /* NACK -> ignore? */
		case 0xff: /* KBD ERROR -> ignore */
			Handled = TRUE;
	}

	if (Handled)
		return 0;

	//
	// Try  the scancode to the corresponding keycode
	//

	if((gKeyboardStatus.lshift || gKeyboardStatus.rshift) && \
	  gScanCodeMap[Char] > 0x2f && gScanCodeMap[Char] < 0x3a) {
		
		//
		// Map special chars above numbers, us-std keymap
	    //

		switch(gScanCodeMap[Char]) {
		case '1':
		  return (UCHAR)'!';
		case '2':
		  return (UCHAR)'@';
		case '3':
		  return (UCHAR)'#';
		case '4':
		  return (UCHAR)'$';
		case '5':
		  return (UCHAR)'%';
		case '6':
		  return (UCHAR)'^';
		case '7':
		  return (UCHAR)'&';
		case '8':
		  return (UCHAR)'*';
		case '9':
		  return (UCHAR)'(';
		case '0':
		  return (UCHAR)')';
		}
	}

	if((gKeyboardStatus.lshift || gKeyboardStatus.rshift) && \
	  gScanCodeMap[Char] < 0x7b && gScanCodeMap[Char] > 0x60)
		return CharToUpper(gScanCodeMap[Char]);
	else
		return gScanCodeMap[Char];
}


VOID 
InitScanCodeMap(
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	RtlZeroMemory(gScanCodeMap, sizeof(gScanCodeMap));

	gScanCodeMap[SCANCODE_KEYPAD0] = '0';
	gScanCodeMap[SCANCODE_KEYPAD1] = '1';
	gScanCodeMap[SCANCODE_KEYPAD2] = '2';
	gScanCodeMap[SCANCODE_KEYPAD3] = '3';
	gScanCodeMap[SCANCODE_KEYPAD4] = '4';
	gScanCodeMap[SCANCODE_KEYPAD5] = '5';
	gScanCodeMap[SCANCODE_KEYPAD6] = '6';
	gScanCodeMap[SCANCODE_KEYPAD7] = '7';
	gScanCodeMap[SCANCODE_KEYPAD8] = '8';
	gScanCodeMap[SCANCODE_KEYPAD9] = '9';
	gScanCodeMap[SCANCODE_KEYPADPLUS]  = '+';
	gScanCodeMap[SCANCODE_KEYPADMINUS] = '-';

	gScanCodeMap[SCANCODE_0] = '0';
	gScanCodeMap[SCANCODE_1] = '1';
	gScanCodeMap[SCANCODE_2] = '2';
	gScanCodeMap[SCANCODE_3] = '3';
	gScanCodeMap[SCANCODE_4] = '4';
	gScanCodeMap[SCANCODE_5] = '5';
	gScanCodeMap[SCANCODE_6] = '6';
	gScanCodeMap[SCANCODE_7] = '7';
	gScanCodeMap[SCANCODE_8] = '8';
	gScanCodeMap[SCANCODE_9] = '9';

	gScanCodeMap[SCANCODE_Q] = 'q';
	gScanCodeMap[SCANCODE_W] = 'w';
	gScanCodeMap[SCANCODE_E] = 'e';
	gScanCodeMap[SCANCODE_R] = 'r';
	gScanCodeMap[SCANCODE_T] = 't';
	gScanCodeMap[SCANCODE_Y] = 'y';
	gScanCodeMap[SCANCODE_U] = 'u';
	gScanCodeMap[SCANCODE_I] = 'i';
	gScanCodeMap[SCANCODE_O] = 'o';
	gScanCodeMap[SCANCODE_P] = 'p';
	gScanCodeMap[SCANCODE_A] = 'a';
	gScanCodeMap[SCANCODE_S] = 's';
	gScanCodeMap[SCANCODE_D] = 'd';
	gScanCodeMap[SCANCODE_F] = 'f';
	gScanCodeMap[SCANCODE_G] = 'g';
	gScanCodeMap[SCANCODE_H] = 'h';
	gScanCodeMap[SCANCODE_J] = 'j';
	gScanCodeMap[SCANCODE_K] = 'k';
	gScanCodeMap[SCANCODE_L] = 'l';
	gScanCodeMap[SCANCODE_Z] = 'z';
	gScanCodeMap[SCANCODE_X] = 'x';
	gScanCodeMap[SCANCODE_C] = 'c';
	gScanCodeMap[SCANCODE_V] = 'v';
	gScanCodeMap[SCANCODE_B] = 'b';
	gScanCodeMap[SCANCODE_N] = 'n';
	gScanCodeMap[SCANCODE_M] = 'm';

	gScanCodeMap[SCANCODE_SPACE]  = ' ';
	gScanCodeMap[SCANCODE_TAB]    = '\t';
	gScanCodeMap[SCANCODE_ENTER]  = '\n';
	gScanCodeMap[SCANCODE_COMMA]  = ',';
	gScanCodeMap[SCANCODE_POINT]  = '.';
	gScanCodeMap[SCANCODE_FSLASH] = '\\';
	gScanCodeMap[SCANCODE_BSLASH] = '/';
	gScanCodeMap[SCANCODE_BACKSPACE] = '\b';
	gScanCodeMap[SCANCODE_EQUAL]  = '=';
}